# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: chensili

import datetime
import json
import logging
import os

from absl import app, flags

from coin.proto.coin_feed_pb2 import SubscriptionLog
from coin.support.proto_log.app.strat_log_replayer import _get_start_end_datetime
from coin.support.proto_log.logic.util import (
    read_subscription_log, SubscriptionInfo)
from coin.util.queue.config import KafkaConfig
from coin.util.queue.constants import TopicType
import google.protobuf.json_format as json_format

FLAGS = flags.FLAGS


class OnLog(object):
  def __init__(self):
    self.record_list = []

  def on_log(self, timestamp, log):
    pb = SubscriptionLog()
    pb.ParseFromString(log)
    data = json_format.MessageToDict(pb, preserving_proto_field_name=True)

    hr_ts = datetime.datetime.fromtimestamp(timestamp / 1e9).strftime(
        '%Y%m%d %H:%M:%S.%f')
    data['record_timestamp'] = timestamp
    data['record_timestamp_human_readable'] = hr_ts
    self.record_list.append(data)

  def on_kafka_log(self, partition, record):
    self.on_log(record.timestamp * 1e6, record.value)


def main(_):
  root_dir = os.path.expanduser(FLAGS.root_dir)
  kafka_config_filename = FLAGS.kafka_config_filename
  trading_date = FLAGS.trading_date
  machine = FLAGS.machine
  market_type = FLAGS.market_type
  exchange = FLAGS.exchange
  api_version = FLAGS.api_version
  start_time = FLAGS.start_time
  end_time = FLAGS.end_time
  assert machine, '--machine must be specified.'
  assert market_type, '--market_type must be specified.'
  assert exchange, '--exchange must be specified.'
  assert api_version, '--api_version must be specified.'

  kafka_config = None
  if kafka_config_filename is not None:
    kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)

  start_dt, end_dt = _get_start_end_datetime(trading_date, start_time, end_time)
  dt_format = '%Y%m%dT%H%M%S'
  start_dt_str = start_dt.strftime(dt_format)
  end_dt_str = end_dt.strftime(dt_format)
  print('Running for %s-%s %s ...' % (start_dt_str, end_dt_str, machine))
  onlog = OnLog()
  callback = onlog.on_kafka_log if kafka_config is not None else onlog.on_log

  sub_info = SubscriptionInfo(market_type=market_type, exchange=exchange,
                              api_version=api_version, machine=machine,
                              trading_date=None)
  kafka_topic_type = TopicType[FLAGS.topic_type] if \
                     FLAGS.topic_type is not None else None
  read_subscription_log(subscription_info=sub_info, start_time=start_dt, end_time=end_dt,
                        callback=callback, root_dir=root_dir,
                        kafka_config=kafka_config,
                        kafka_topic_type=kafka_topic_type)
  filename = 'subscription_log.%s-%s.%s.%s.%s.json' % (
      start_dt_str, end_dt_str, market_type, exchange, api_version)
  with open(filename, 'w') as out_file:
    if FLAGS.unpretty:
      records = '[\n' + ',\n'.join(json.dumps(record) for record in onlog.record_list) + '\n]'
      out_file.write(records)
    else:
      json.dump(onlog.record_list, out_file, indent=2)
  print(filename)


if __name__ == '__main__':
  flags.DEFINE_string('root_dir', '~/data/ticker_proto_log', 'root_dir.')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('topic_type', None, 'topic type')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('market_type', None, 'market_type')
  flags.DEFINE_string('exchange', None, 'exchange')
  flags.DEFINE_string('api_version', None, 'api_version')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('unpretty', False, '')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)

